#include <iostream>
#include <vector>
#include <string>
#include <math.h>
#include <fstream>
#include <sstream>
using namespace std;

// RETURNS RANDOM NUMBER IN 0 TO 1
double randFloat(void){
    return ((double) rand())/(double)RAND_MAX;
}

// RETURNS WHETHER ANY NON-TARGET VALUES EXIST AT SUPPLIED LEVEL
bool badAt(vector<int> X, vector<bool> S, int level){
    for(int i=0;i<X.size();i++){
        if(X[i]==level){
            if(!S[i]){
                return true;
            }
        }
    }
    return false;
}

// LEVEL C BIT-FLIPPING (RETURNS TIME TAKEN FOR ALL 1)
double modifyC(vector<int> X,int TC){
    int countC = 0;
    for(int i=0;i<X.size();i++){
        if(X[i]==2){
            countC++;
        }
    }
    for(int tC=0;tC<TC;tC++){
        bool allOne = true;
        for(int i=0;i<countC;i++){
            if(randFloat()<0.5){
                allOne=false;
            }
        }
        if(allOne){
            return (tC+1)/(double)TC;
        }
    }
    return 1.0;
}

// LEVEL B MODS
void modifyB(vector<int> &X,vector<bool> &S, double prob){
    for(int i=0;i<X.size();i++){
        if(X[i]==1){
            if(randFloat()<prob){
                S[i]=!S[i];
            }
        }
    }
}

// LEVEL A MODS
void modifyA(vector<int> &X,vector<bool> &S, double prob){
    for(int i=0;i<X.size();i++){
        if(randFloat()<prob){
            X[i] = floor(randFloat()*3);
            S[i] = (randFloat()<0.5);
        }
    }
}



int main (int argc, char **argv){

    // MODEL PARAMS
    int N = 20;
    int TA = 1000;
    int TB = 5000;
    int TC = 5000;
    double k = 50.;

    std::stringstream ss;
    ofstream logfile;

    // OVERWRITE DEFAULT PARAMS WITH COMMAND LINE ARGS AND WRITE PARAMS TO FILE
    if(argc>1){
        ss<<argv[1];
        logfile.open(ss.str().c_str(),ios::out|ios::app);
    } else {
        cout<<"Supply filename for logs."<<endl;
        return 0;
    }
    if(argc>2){
        srand(stoi(argv[2]));
        logfile<<"Seed: "<<argv[2]<<endl;
    } else {
        logfile<<"Seed: none"<<endl;
    }
    if(argc>3){ N = stoi(argv[3]); }
    if(argc>4){ TA = stoi(argv[4]); }
    if(argc>5){ TB = stoi(argv[5]); }
    if(argc>6){ TC = stoi(argv[6]); }
    if(argc>7){ k = stod(argv[7]); }

    logfile<<"N: "<<N<<endl;
    logfile<<"TA: "<<TA<<endl;
    logfile<<"TB: "<<TB<<endl;
    logfile<<"TC: "<<TC<<endl;
    logfile<<"k: "<<k<<endl;

    // FOR RECORDING TARGET DISCOVERY TIME
    bool firstFound = false;
    int foundTA=0;
    int foundTB=0;

    // INITIAL CONDITIONS
    double pAmod = 0.5;
    double pBmod = 0.5;
    vector<int> X(N);   // LEVEL OF EACH CHAR (0=C, 1=B, 2=A)
    vector<bool> S(N);  // BINARY STATE OF EACH CHAR
    for(int i=0;i<N;i++){
        X[i] = floor(randFloat()*3);
        S[i] = (randFloat()<0.5);
    }

    // STORAGE FOR POTENTIAL RESET
    vector<int> XcpA(N);
    vector<int> XcpB(N);
    vector<bool> ScpA(N);
    vector<bool> ScpB(N);

    // LEVEL A LOOP
    double tauBp = 1.;
    for(int tA=0;tA<TA;tA++){
        XcpA = X;
        ScpA = S;

        modifyA(X,S,pAmod);

        // LEVEL B EVAL
        double tauB = 1.;
        if(!badAt(X,S,0)){
            // LEVEL B LOOP
            double tauCp = 1.;
            for(int tB=0;tB<TB;tB++){
                XcpB = X;
                ScpB = S;

                modifyB(X,S,pBmod);

                // LEVEL C EVAL
                double tauC = 1.;

                if(!badAt(X,S,1)){

                    tauC = modifyC(X,TC);

                    if(tauC<1.){
                       tauB = (tB+1)/(double)TB;
                       if(!firstFound){
                            foundTB = tB;
                            foundTA = tA;
                            firstFound = true;
                        }
                    }
                }

                // LEVEL B MODIFY
                pBmod = tauC*0.5;
                double pBapt = 1./(1+exp(k*(tauC-tauCp)));
                if(randFloat()<pBapt){
                    tauCp = tauC;
                } else {
                    X = XcpB;
                    S = ScpB;
                    tauC = tauCp;
                }
            }
        }

        // LEVEL A MODIFY
        pAmod = tauB*0.5;
        double pAapt = 1./(1+exp(k*(tauB-tauBp)));
        if(randFloat()<pAapt){
            tauBp = tauB;
        } else {
            X = XcpA;
            S = ScpA;
            tauB = tauBp;
        }

    }


    // PRINTOUT
    int countA=0;
    int countB=0;
    int countC=0;
    for(int i=0;i<N;i++){
        if(X[i]==0){
            logfile<<"A,";
            countA++;
        } else if(X[i]==1){
            logfile<<"B,";
            countB++;
        } else {
            logfile<<"C,";
            countC++;
        }
    }
    logfile<<endl;
    for(int i=0;i<N;i++){
        if(X[i]==2){
            logfile<<"?,";
        } else {
            logfile<<S[i]<<",";
        }
    }
    logfile<<endl;
    if(badAt(X,S,0) || badAt(X,S,1)){
        logfile<<"BAD"<<endl;
    } else {
        logfile<<"GOOD"<<endl;
    }
    logfile<<"proportion at level A: "<<(double)countA/(double)N<<endl;
    logfile<<"proportion at level B: "<<(double)countB/(double)N<<endl;
    logfile<<"proportion at level C: "<<(double)countC/(double)N<<endl;
    logfile<<"target found at T_A: "<<foundTA<<endl;
    logfile<<"target found at T_B: "<<foundTB<<endl;

    logfile.close();

    return 0;
}
